# examples/pytorch_model_utils.py
"""
PyTorch模型加载和推理工具
提供模型保存、加载、推理等实用功能
"""

import torch
import torch.nn as nn
import numpy as np
import pandas as pd
from pathlib import Path
from typing import Dict, Any, List
import joblib

class ModelManager:
    """PyTorch模型管理器"""
    
    def __init__(self, models_dir: str = "models"):
        self.models_dir = Path(models_dir)
        self.models_dir.mkdir(exist_ok=True)
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    def save_complete_model(self, model: nn.Module, preprocessor: Any, 
                          model_name: str, metadata: Dict = None):
        """保存完整模型（包括预处理器和元数据）"""
        
        save_path = self.models_dir / f"{model_name}_complete.pth"
        
        # 准备保存数据
        save_data = {
            'model_state_dict': model.state_dict(),
            'model_class': model.__class__.__name__,
            'preprocessor': preprocessor,
            'device': str(self.device),
            'timestamp': pd.Timestamp.now().isoformat()
        }
        
        # 添加元数据
        if metadata:
            save_data['metadata'] = metadata
        
        # 保存模型
        torch.save(save_data, save_path)
        
        # 单独保存预处理器（使用joblib）
        joblib.dump(preprocessor, self.models_dir / f"{model_name}_preprocessor.pkl")
        
        print(f"✅ 完整模型已保存到: {save_path}")
        
        return save_path
    
    def load_complete_model(self, model_name: str, model_class: nn.Module):
        """加载完整模型"""
        
        save_path = self.models_dir / f"{model_name}_complete.pth"
        
        if not save_path.exists():
            raise FileNotFoundError(f"模型文件不存在: {save_path}")
        
        # 加载模型数据
        checkpoint = torch.load(save_path, map_location=self.device)
        
        # 重建模型
        model = model_class
        model.load_state_dict(checkpoint['model_state_dict'])
        model.to(self.device)
        model.eval()
        
        # 加载预处理器
        preprocessor_path = self.models_dir / f"{model_name}_preprocessor.pkl"
        preprocessor = joblib.load(preprocessor_path)
        
        print(f"✅ 模型已加载: {model_name}")
        print(f"   设备: {checkpoint.get('device', 'unknown')}")
        print(f"   保存时间: {checkpoint.get('timestamp', 'unknown')}")
        
        return model, preprocessor, checkpoint.get('metadata', {})
    
    def list_saved_models(self) -> List[str]:
        """列出已保存的模型"""
        
        model_files = list(self.models_dir.glob("*_complete.pth"))
        model_names = [f.stem.replace('_complete', '') for f in model_files]
        
        return model_names

class RealTimePredictor:
    """实时预测器"""
    
    def __init__(self, model: nn.Module, preprocessor: Any, device: torch.device):
        self.model = model
        self.preprocessor = preprocessor
        self.device = device
        
        # 预测历史
        self.prediction_history = []
        
    def predict_single(self, data: pd.DataFrame) -> Dict:
        """单次预测"""
        
        try:
            # 特征工程
            processed_data = self.preprocessor.create_features(data)
            
            # 创建序列
            X, _ = self.preprocessor.create_sequences(processed_data, target_column='close')
            
            if len(X) == 0:
                raise ValueError("数据不足以创建预测序列")
            
            # 取最后一个序列进行预测
            last_sequence = X[-1]
            sequence_tensor = torch.FloatTensor(last_sequence).unsqueeze(0).to(self.device)
            
            # 模型预测
            self.model.eval()
            with torch.no_grad():
                prediction = self.model(sequence_tensor).cpu().numpy()[0, 0]
            
            # 反向转换
            pred_price = self.preprocessor.inverse_transform_predictions(np.array([prediction]))[0]
            
            # 计算变化
            current_price = data['close'].iloc[-1]
            change_pct = ((pred_price - current_price) / current_price) * 100
            
            # 保存预测历史
            prediction_result = {
                'timestamp': pd.Timestamp.now(),
                'current_price': current_price,
                'predicted_price': pred_price,
                'change_percent': change_pct,
                'raw_prediction': prediction
            }
            
            self.prediction_history.append(prediction_result)
            
            return prediction_result
            
        except Exception as e:
            print(f"❌ 预测失败: {e}")
            return None
    
    def predict_batch(self, data_list: List[pd.DataFrame]) -> List[Dict]:
        """批量预测"""
        
        results = []
        
        for i, data in enumerate(data_list):
            print(f"🔮 预测 {i+1}/{len(data_list)}...")
            result = self.predict_single(data)
            results.append(result)
        
        return results
    
    def get_prediction_statistics(self) -> Dict:
        """获取预测统计信息"""
        
        if not self.prediction_history:
            return {}
        
        changes = [pred['change_percent'] for pred in self.prediction_history]
        
        return {
            'total_predictions': len(self.prediction_history),
            'avg_change_percent': np.mean(changes),
            'std_change_percent': np.std(changes),
            'min_change_percent': np.min(changes),
            'max_change_percent': np.max(changes),
            'latest_prediction': self.prediction_history[-1]
        }

class ModelEvaluator:
    """模型评估器"""
    
    def __init__(self, device: torch.device):
        self.device = device
    
    def evaluate_model_comprehensive(self, model: nn.Module, test_loader, 
                                   preprocessor: Any, symbol: str) -> Dict:
        """综合评估模型"""
        
        model.eval()
        predictions = []
        targets = []
        
        # 获取预测和真实值
        with torch.no_grad():
            for batch_x, batch_y in test_loader:
                batch_x, batch_y = batch_x.to(self.device), batch_y.to(self.device)
                outputs = model(batch_x)
                predictions.extend(outputs.squeeze().cpu().numpy())
                targets.extend(batch_y.cpu().numpy())
        
        predictions = np.array(predictions)
        targets = np.array(targets)
        
        # 反向转换到原始价格空间
        pred_prices = preprocessor.inverse_transform_predictions(predictions)
        target_prices = preprocessor.inverse_transform_predictions(targets)
        
        # 计算各种指标
        mse = np.mean((predictions - targets) ** 2)
        mae = np.mean(np.abs(predictions - targets))
        rmse = np.sqrt(mse)
        
        # 价格空间的指标
        price_mse = np.mean((pred_prices - target_prices) ** 2)
        price_mae = np.mean(np.abs(pred_prices - target_prices))
        price_rmse = np.sqrt(price_mse)
        
        # 方向准确率
        if len(targets) > 1:
            actual_direction = np.diff(targets) > 0
            pred_direction = np.diff(predictions) > 0
            direction_accuracy = np.mean(actual_direction == pred_direction)
        else:
            direction_accuracy = 0
        
        # 相关系数
        correlation = np.corrcoef(predictions, targets)[0, 1]
        
        # MAPE (Mean Absolute Percentage Error)
        mape = np.mean(np.abs((target_prices - pred_prices) / target_prices)) * 100
        
        # 评估结果
        evaluation_results = {
            'symbol': symbol,
            'sample_size': len(predictions),
            'normalized_metrics': {
                'mse': mse,
                'mae': mae,
                'rmse': rmse
            },
            'price_metrics': {
                'mse': price_mse,
                'mae': price_mae,
                'rmse': price_rmse,
                'mape': mape
            },
            'direction_accuracy': direction_accuracy,
            'correlation': correlation,
            'predictions_sample': pred_prices[:10].tolist(),
            'targets_sample': target_prices[:10].tolist()
        }
        
        return evaluation_results
    
    def generate_evaluation_report(self, evaluation_results: Dict) -> str:
        """生成评估报告"""
        
        report = f"""
📊 模型评估报告
{'='*50}
股票代码: {evaluation_results['symbol']}
样本数量: {evaluation_results['sample_size']}

📈 标准化指标:
   MSE: {evaluation_results['normalized_metrics']['mse']:.6f}
   MAE: {evaluation_results['normalized_metrics']['mae']:.6f}
   RMSE: {evaluation_results['normalized_metrics']['rmse']:.6f}

💰 价格指标:
   价格MSE: ${evaluation_results['price_metrics']['mse']:.2f}
   价格MAE: ${evaluation_results['price_metrics']['mae']:.2f}
   价格RMSE: ${evaluation_results['price_metrics']['rmse']:.2f}
   MAPE: {evaluation_results['price_metrics']['mape']:.2f}%

🎯 其他指标:
   方向准确率: {evaluation_results['direction_accuracy']:.1%}
   相关系数: {evaluation_results['correlation']:.3f}

📋 样本预测 (前10个):
"""
        
        for i in range(min(10, len(evaluation_results['predictions_sample']))):
            pred = evaluation_results['predictions_sample'][i]
            target = evaluation_results['targets_sample'][i]
            error = abs(pred - target)
            report += f"   {i+1:2d}. 预测: ${pred:7.2f}, 实际: ${target:7.2f}, 误差: ${error:6.2f}\n"
        
        # 评级
        if evaluation_results['direction_accuracy'] > 0.6:
            rating = "优秀 ⭐⭐⭐⭐⭐"
        elif evaluation_results['direction_accuracy'] > 0.55:
            rating = "良好 ⭐⭐⭐⭐"
        elif evaluation_results['direction_accuracy'] > 0.52:
            rating = "一般 ⭐⭐⭐"
        else:
            rating = "需改进 ⭐⭐"
        
        report += f"\n🏆 模型评级: {rating}\n"
        
        return report

def demo_pytorch_model_usage():
    """演示PyTorch模型使用流程"""
    
    print("🚀 PyTorch模型使用演示")
    print("=" * 50)
    
    # 1. 创建模型管理器
    manager = ModelManager()
    
    # 2. 列出已保存的模型
    saved_models = manager.list_saved_models()
    print(f"📂 已保存的模型: {saved_models}")
    
    # 3. 模拟实时预测流程
    print(f"\n🔮 实时预测演示:")
    print(f"   (这里需要实际的训练好的模型)")
    
    # 4. 模型评估演示
    print(f"\n📊 模型评估演示:")
    print(f"   (这里需要实际的测试数据)")
    
    # 5. 使用建议
    print(f"\n💡 PyTorch模型使用建议:")
    print(f"   1. 🔧 定期重新训练模型")
    print(f"   2. 📊 监控模型性能")
    print(f"   3. 🎯 收集反馈数据")
    print(f"   4. ⚡ 优化推理速度")
    print(f"   5. 🛡️ 设置异常处理")

if __name__ == "__main__":
    demo_pytorch_model_usage()
